import pdb
from numpy import  *
import string

"""
Optimal-Marker-Panel (OPM) is an algorithm for cost-efficient selection of a
marker panel in genetic studies.

Please see associated ReadMe file for data input description.

Open-Source License: The MIT License

Copyright (c) 2011 Jamie S. Sanderlin, Nicole Lazar, Michael J. Conroy, and Jaxk Reeves

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE."""

def PIDsib(props):
    """Computes the probability of identity among siblings 

    props: vector of allele frequencies
    
    Returns the probability of identity among siblings at one locus."""
    
    square=[]
    four=[]
    
    for i in range(len(props)):
        square.append(props[i]**2)
        four.append(props[i]**4)
        
    prob = 0.25 + (0.5*float(sum(square))) + (0.5*float(sum(square)**2)) - (0.25*float(sum(four)))
    return prob

def Pex(props):
    """Computes the probability of exclusion (Jamieson 1979,1994)

    props: vector of allele frequencies
    
    Returns the probability of exclusion at one locus."""

    square=[]
    cube=[]
    four=[]
    five=[]
    
    for i in range(len(props)):
        square.append(props[i]**2)
        cube.append(props[i]**3)
        four.append(props[i]**4)
        five.append(props[i]**5)
    
    prob=1.-2*sum(square)+sum(cube)+2*sum(four)-3*sum(five)-2*((sum(square))**2)+3*sum(square)*sum(cube)

    return prob

def genComb(x, k):
    """Generate all combinations of k elements from list x.

    k: number of elements in one combination
    x: list of all possible elements
    
    Returns a list of all combinations with k elements."""

    if k > len(x):
        return
    if k == 0:
        yield []
    else:
        # A combination does or doesn't contain first.
        first, rest = x[0], x[1:]
        # If it does, the remainder is a k-1 comb of rest.
        for c in genComb(rest, k-1):
            c.insert(0, first)
            yield c
        # If it doesn't contain first, it's a k comb of rest.
        for c in genComb(rest, k):
            yield c

def readData(name1,name2):
    """"Read text files in for genotype frequency and genotyping error

    name1: text file name for genotype frequency data
    name2: text file name for genotyping error data
    
    Returns dictionary of loci names to loci numbers, dictionary of loci
    numbers to genotype frequencies, and dictionary of loci numbers to
    genotyping errors."""

    data1=open(name1,'rb') #genotype frequency data
    text1=data1.readlines()
    
    locinames={}    #dictionary of loci names
    loci={}         #dictionary of loci:genotype frequencies
    loc=0

    for line in text1:    
        clean=line.strip('\r\n').split('\t')
        locname=clean[0]

        if locname in locinames:
            loci[locinames[locname]].append(float(clean[2]))
        else:
            locinames[locname]=loc
            loci[locinames[locname]]=[float(clean[2])]
            loc+=1
          
    data1.close()

    data2=open(name2,'rb')  #genotyping error data (FA,ADO)
    text2=data2.readlines()

    generr={}   #dictionary of gen error [FA,ADO]

    for line in text2:
        clean=line.strip('\r\n').split('\t')
        locname=clean[0]

        generr[locinames[locname]]=[float(clean[1]),float(clean[2])]

    return locinames,loci,generr

def optInput(freqs,generr,par):
    """Combines data to create input for optimization procedure

    freqs: dictionary of loci numbers to genotype frequencies, and 
    generr: dictionary of loci numbers to genotyping errors [FA,ADO]
    par: binary indicator of probability of identity among siblings
         (0) or probability of exclusion (1)

    Returns a dictionary of each locus associated with a vector
    including genotyping error from false alleles, genotyping
    error from allelic dropout, and either probability of identity
    among siblings or probability of exclusion."""

    if par==0:  #PIDsib
        out={}
        for loc in freqs:
            sib=PIDsib(freqs[loc])
            out[loc]=[generr[loc][0],generr[loc][1],sib]
                   
    elif par==1:    #Pex
        out={}
        for loc in freqs:
            ex=Pex(freqs[loc])
            out[loc]=[generr[loc][0],generr[loc][1],ex]
       
    return out

def lociName(locdiction):
    """Creates a reverse dictionary for locus numbers and locus names

    lociName: dictionary of locus names to locus numbers

    Returns dictionary of locus number to locus names."""

    loc=len(locdiction)

    locNumName={}

    for i in range(loc):
        for item in locdiction:
            if locdiction[item]==i:
                locNumName[i]=item

    return locNumName

def optimize(loci,oparms,constrain,opar,cost):
    """Optimize over given parameters and constraints

    loci: dictionary of locus names to locus numbers
    oparms: dictionary of each locus associated with a vector
            including genotyping error from false alleles,
            genotyping error from allelic dropout, and either
            probability of identity among siblings or probability
            of exclusion
    constrain: vector of constaints for mean FA error, mean ADO
               error,probability of identity among siblings or probability
               of exclusion, and either number of loci or cost
    opar: binary indicator of probability of identity among siblings
          (0) or probability of exclusion (1)
    cost: either 'NA' if optimizing over number of loci or
          a vector of initial cost and cost per locus

    Creates output csv file named optimal_data."""

    seq = range(0, len(loci))

    locnamedict=lociName(loci)
    locnames=[locnamedict[k] for k in seq]
    locnames2=string.join(locnames,',')

    sumfile=open('optimal_data.csv','w')
    if opar==0:
        titles=string.join([locnames2,'LociNum','Cost','MFA','MDO','PIDsib','\n'],',')
    else:
        titles=string.join([locnames2,'LociNum','Cost','MFA','MDO','Pex','\n'],',')
    sumfile.writelines(titles)

    comb=[]
    for k in range(len(seq)):
        for c in genComb(seq, k+1): #all combinations with at least 1 locus
            comb.append(c)

    if opar==0:
        for i in range(len(comb)):
            x=[0.]*len(seq)
            for j in range(len(comb[i])):
                x[comb[i][j]]=1. 

            pid=1.
            ar=0.
            fr=0.
            for k in range(len(x)):
                if x[k]==1.:
                    fr+=oparms[k][0]
                    ar+=oparms[k][1]               
                    pid*=oparms[k][2]

            locnum=sum(x)

            if cost=='NA':
                locOrcost=locnum
            else:    
                locOrcost=cost[0]+cost[1]*locnum
                
            if locOrcost<=constrain[3]: #number of loci or cost
 
                if pid<=constrain[2]: #PID constraint        
                    adoerr=ar/sum(x)

                    if adoerr<=constrain[1]: #MDO constraint
                        faerr=fr/sum(x)

                        if faerr<=constrain[0]: #MFA constraint                                                        
                            alldat=[]
                            xstring=[str(x[item]) for item in range(len(x))]
                            xstring2=string.join(xstring,',')
                            if cost=='NA':
                                dat=string.join([xstring2,str(locnum),cost,str(faerr),str(adoerr),str(pid)],',')
                            else:
                                dat=string.join([xstring2,str(locnum),str(locOrcost),str(faerr),str(adoerr),str(pid)],',') 
                            alldat.append(dat)
                            alldat.append("\n")
                            
                            sumfile.writelines(alldat)
   
        sumfile.close()

    else:    

        for i in range(len(comb)):
            x=[0.]*len(seq)
            for j in range(len(comb[i])):
                x[comb[i][j]]=1.

            pex=1.     
            ar=0.
            fr=0.
            for k in range(len(x)):
                if x[k]==1.:
                    fr+=oparms[k][0]
                    ar+=oparms[k][1]
                    pex*=(1.-oparms[k][2])

            pextot=1.-pex

            locnum=sum(x)
            
            if cost=='NA':
                locOrcost=locnum
            else:    
                locOrcost=cost[0]+cost[1]*locnum
                
            if locOrcost<=constrain[3]:  #number of loci or cost

                if pextot>=constrain[2]: #Pex constraint        
                    adoerr=ar/sum(x)

                    if adoerr<=constrain[1]: #MDO constraint
                        faerr=fr/sum(x)

                        if faerr<=constrain[0]: #MFA constraint                                                        
                            alldat=[]
                            xstring=[str(x[item]) for item in range(len(x))]
                            xstring2=string.join(xstring,',')
                            if cost=='NA':
                                dat=string.join([xstring2,str(locnum),cost,str(faerr),str(adoerr),str(pextot)],',')
                            else:
                                dat=string.join([xstring2,str(locnum),str(locOrcost),str(faerr),str(adoerr),str(pextot)],',')
                            alldat.append(dat)
                            alldat.append("\n")
                            
                            sumfile.writelines(alldat)
   
        sumfile.close()

def isInteger(number):
    """Checks number if it is an integer of either 0 or 1

    Returns Boolean."""

    if not isinstance(number,int):
        raise TypeError, "An integer value is required."
        return False

    if not (number==0 or number==1):
        raise TypeError, "An integer equal to 0 or 1 is required."
        return False

    return True

def isIntegerLoci(number):
    """Checks number if it is an integer greater than or equal to 1

    Returns Boolean."""

    if not isinstance(number,int):
        raise TypeError, "An integer value is required."
        return False

    if not (number>=1):
        raise TypeError, "An integer greater than 0 is required."
        return False

    return True

def isFloatProb(number):
    """Checks number if it is a real number bound between 0 and 1

    Returns Boolean."""

    if not isinstance(number,float):
        raise TypeError, "An real value is required."
        return False

    if not (number>=0 and number<=1):
        raise TypeError, "A real value bound between 0 and 1 is required."
        return False

    return True

def isFloatCost(number):
    """Checks number if it is a real number > 0

    Returns Boolean."""

    if not isinstance(number,float):
        raise TypeError, "An real value is required."
        return False

    if not (number>=0):
        raise TypeError, "A real value greater than 0.0 is required."
        return False

    return True
    
def main():
    """User inputs data file names, constraints, PID or Pex"""

    #add more user-input validation here
    
    file1=raw_input("Enter genotype frequency text file name:") 
    file2=raw_input("Enter genotyping error text file name:")

    parm=input("Optimize with PIDsib (enter 0) or Pex (enter 1):")
    parmtest=isInteger(parm)

    costloci=input("Optimize over loci (enter 0) or cost (enter 1):")
    costlocitest=isInteger(costloci)

    cMFA=input("Enter mean expected false allele error constraint:")
    cMFAtest=isFloatProb(cMFA)

    cMDO=input("Enter mean expected dropout error constraint:")
    cMDOtest=isFloatProb(cMDO)

    if costloci==0:
        cLoci=input("Enter loci constraint:")
        clocitest=isIntegerLoci(cLoci)
    else:
        cCost=input("Enter cost constraint as real number:")
        ccost1=isFloatCost(cCost)
        dollars=[0.0]*2
        dollars[0]=input("Enter initial cost as real number:")
        ccost2=isFloatCost(dollars[0])
        dollars[1]=input("Enter per locus cost as real number:")
        ccost3=isFloatCost(dollars[1])

    if parm==0:
        cPIDsib=input("Enter PIDsib constraint:")
        cPIDsibtest=isFloatProb(cPIDsib)
        
        lnames,freq,gerr=readData(file1+'.txt',file2+'.txt')
        opt=optInput(freq,gerr,parm)

        if costloci==0:
            constraints=[cMFA,cMDO,cPIDsib,cLoci]
            optimize(lnames,opt,constraints,parm,'NA')
        else:
            constraints=[cMFA,cMDO,cPIDsib,cCost]
            optimize(lnames,opt,constraints,parm,dollars)
            
    else:
        cPex=input("Enter Pex constraint:")
        cPextest=isFloatProb(cPex)
        
        lnames,freq,gerr=readData(file1+'.txt',file2+'.txt')
        opt=optInput(freq,gerr,parm)

        if costloci==0:      
            constraints=[cMFA,cMDO,cPex,cLoci]
            optimize(lnames,opt,constraints,parm,'NA')
        else:
            constraints=[cMFA,cMDO,cPex,cCost]
            optimize(lnames,opt,constraints,parm,dollars)

    print "An optimization file has been created."

main()
